# -*- coding: utf-8 -*-
import os
import random
import time

import redis

from setting.mysql_test import mysql45_config
from embed_text.embed_gpt_text import GPTCall
import pymysql
from utils.msg_queue import connect_message_queue
from setting.redis_config import redis_embed_url
# os.environ['HTTP_PROXY'] = 'http://192.168.1.45:6666'
# os.environ['HTTPS_PROXY'] = 'https://192.168.1.45:6666'
REDIS_PASSWORD = None


class EmbedDataProduce(object):
    def __init__(self):
        self.mysql_conn = pymysql.connect(**mysql45_config)
        self.mysql_cursor = self.mysql_conn.cursor()
        self.redis_url = redis_embed_url
        self.redis_conn = redis.Redis(host='localhost', port=6379, db=3, password=REDIS_PASSWORD)

    def get_tag_tittle(self):
        #sql = 'SELECT id,tag_title  FROM image_info WHERE download_status = 1 AND width / height BETWEEN 1.38 AND 1.7 AND embeddings_labels IS NULL'
        sql = "SELECT id,tag_title  FROM image_info WHERE download_status = 1 AND embeddings_labels IS NULL"
        #sql = "SELECT id,tag_title  FROM image_info WHERE source= 'local'"
        self.mysql_cursor.execute(sql)
        tag_data = self.mysql_cursor.fetchall()
        return tag_data

    def put_redis(self, queue_name, json_msg):
        #避免openAI的接口不稳定导致的数据丢失，将数据存入redis，控制速度如果失败重新入队
        #生产写入数据
        p = connect_message_queue(queue_name, url=self.redis_url, maxsize=10000, lazy_limit=True)
        p.put(json_msg)

    def run(self,queue_name):
        gptcall = GPTCall()
        tag_data = self.get_tag_tittle()
        if not tag_data:
            print('tag_data is empty')
            return
        print('tag_data len:', len(tag_data))
        for i in tag_data:
            image_id = i[0]
            tag = i[1]
            print('tag:',tag)
            #判断图片是否已经成功获得embed，如果成功则跳过
            if self.redis_conn.sismember('embed_label_set', image_id):
                print('图片id已存在，已经解析成功，跳过')
                continue
            try:
                tag_embedding = gptcall.chatgpt_azure_embeding(tag)
                #tag_embedding = gptcall.other_azure_embeding(tag)
                self.redis_conn.sadd('embed_label_set', image_id)
                time.sleep(random.randint(1, 2))
            except Exception as e:
                print('embed api error:',e)
                retry_queue_name = queue_name + '_retry'
                msg = {'id': image_id, 'tag_title': tag}
                self.put_redis(retry_queue_name, msg)
                continue

            json_msg = {'id': image_id, 'tag_title': tag, 'tag_embedding': tag_embedding}
            self.put_redis(queue_name, json_msg)

    def retry_run(self,retry_queue_name):
        gptcall = GPTCall()
        q = connect_message_queue(retry_queue_name, url=self.redis_url, maxsize=10000, lazy_limit=True)
        while q.qsize() > 0:
            json_msg = q.get()
            print('json_msg', json_msg)
            if not json_msg:
                continue
            id = json_msg['id']
            tag = json_msg['tag_title']
            # 判断图片是否已经成功获得embed，如果成功则跳过
            if self.redis_conn.sismember('embed_label_set', id):
                print('图片id已存在，已经解析成功，跳过')
                continue
            try:
                tag_embedding = gptcall.chatgpt_azure_embeding(tag)
                # tag_embedding = gptcall.other_azure_embeding(tag)
                self.redis_conn.sadd('embed_label_set', id)
                time.sleep(random.randint(1, 2))
            except Exception as e:
                print('embed api error:',e)
                #second_retry_queue_name = retry_queue_name + '_retry'
                second_retry_queue_name = retry_queue_name
                msg = {'id': id, 'tag_title': tag}
                self.put_redis(second_retry_queue_name, msg)
                continue

            json_msg = {'id': id, 'tag_title': tag, 'tag_embedding': tag_embedding}
            self.put_redis(queue_name, json_msg)


if __name__ == '__main__':
    queue_name = 'tag_embed_data'
    produce = EmbedDataProduce()
    # produce.run(queue_name)
    retry_queue_name = queue_name + '_retry'
    produce.retry_run(retry_queue_name)


